﻿namespace Microsoft.VisualStudio.BaseCustomTool
{
    using System;
    using System.IO;
    using System.Linq;
    using System.Runtime.InteropServices;
    using System.Text;
    using Microsoft.Win32;

    public abstract class BaseCustomTool : IVsSingleFileGenerator
    {
        protected BaseCustomTool()
        {
            this.InputFilePath = string.Empty;
            this.FileNamespace = string.Empty;
        }

        public void Generate(string wszInputFilePath, string bstrInputFileContents, string wszDefaultNamespace, out IntPtr rgbOutputFileContents, out int pcbOutput, IVsGeneratorProgress pGenerateProgress)
        {
            if (bstrInputFileContents == null)
            {
                throw new ArgumentNullException(bstrInputFileContents);
            }
            this.InputFilePath = wszInputFilePath;
            this.FileNamespace = wszDefaultNamespace;
            this.CodeGeneratorProgress = pGenerateProgress;
            byte[] source = this.GenerateCode(wszInputFilePath, bstrInputFileContents);
            if (source == null)
            {
                rgbOutputFileContents = IntPtr.Zero;
                pcbOutput = 0;
            }
            else
            {
                pcbOutput = source.Length;
                rgbOutputFileContents = Marshal.AllocCoTaskMem(pcbOutput);
                Marshal.Copy(source, 0, rgbOutputFileContents, pcbOutput);
            }
        }

        protected abstract byte[] GenerateCode(string inputFileName, string inputFileContent);

        protected virtual void GeneratorErrorCallback(bool warning, int level, string message, int line, int column)
        {
            IVsGeneratorProgress codeGeneratorProgress = this.CodeGeneratorProgress;
            if (codeGeneratorProgress != null)
            {
                codeGeneratorProgress.GeneratorError(warning, level, message, line, column);
            }
        }

        public abstract string GetDefaultExtension();

        protected static byte[] StreamToBytes(Stream stream)
        {
            if (stream.Length == 0L)
            {
                return new byte[0];
            }
            long position = stream.Position;
            stream.Position = 0L;
            byte[] buffer = new byte[(int)stream.Length];
            stream.Read(buffer, 0, buffer.Length);
            stream.Position = position;
            return buffer;
        }

        #region COM Registration

        [ComRegisterFunction]
        public static void ComRegister(Type type)
        {
            var attributes = Attribute.GetCustomAttributes(type, typeof(CustomToolRegistrationAttribute), false).Cast<CustomToolRegistrationAttribute>();
            foreach (var attr in attributes)
            {
                try
                {
                    var generatorsKeyName = GetVsGeneratorsRoot(attr.VsVersion, attr.PrjKind);
                    var generatorsKey = Registry.LocalMachine.OpenSubKey(generatorsKeyName, true);

                    var customToolKey = generatorsKey.CreateSubKey(attr.Name);
                    customToolKey.SetValue(string.Empty, attr.Description, RegistryValueKind.String);
                    customToolKey.SetValue("CLSID", type.GUID.ToString("B"), RegistryValueKind.String);
                    customToolKey.SetValue("GeneratesDesignTimeSource", attr.GeneratesDesignTimeSource, RegistryValueKind.DWord);
                    Console.WriteLine("Successfully registered custom tool for Visual Studio version {0} for project kind {1}", attr.VsVersion, attr.PrjKind);
                }
                catch (Exception ex)
                {
                    Console.WriteLine("Error registering custom tool for Visual Studio version {0} for project kind {1}. Exception details:\n{2}", attr.VsVersion, attr.PrjKind, ex.ToString());
                }
            }
        }

        [ComUnregisterFunction]
        public static void ComUnregister(Type type)
        {
            var attributes = Attribute.GetCustomAttributes(type, typeof(CustomToolRegistrationAttribute), false).Cast<CustomToolRegistrationAttribute>();
            foreach (var attr in attributes)
            {
                try
                {
                    var generatorsKeyName = GetVsGeneratorsRoot(attr.VsVersion, attr.PrjKind);
                    var generatorsKey = Registry.LocalMachine.OpenSubKey(generatorsKeyName, true);

                    generatorsKey.DeleteSubKey(attr.Name);
                    Console.WriteLine("Successfully unregistered custom tool for Visual Studio version {0} for project kind {1}", attr.VsVersion, attr.PrjKind);
                }
                catch (Exception ex)
                {
                    Console.WriteLine("Error registering custom tool for Visual Studio version {0} for project kind {1}. Exception details:\n{2}", attr.VsVersion, attr.PrjKind, ex.ToString());
                }
            }
        }

        private static string GetVsGeneratorsRoot(string vsVersion, string prjKind)
        {
            StringBuilder sb = new StringBuilder();

            sb.Append(@"SOFTWARE");
            if (EnvironmentExtensions.Is64BitOperatingSystem)
            {
                sb.Append(@"\Wow6432Node");
            }
            sb.Append(@"\Microsoft\VisualStudio");
            sb.Append(@"\").Append(vsVersion);
            sb.Append(@"\Generators");
            sb.Append(@"\").Append(prjKind);

            return sb.ToString();
        }

        #endregion

        internal IVsGeneratorProgress CodeGeneratorProgress { get; private set; }
        protected string FileNamespace { get; private set; }
        protected string InputFilePath { get; private set; }
    }
}

